import os
import torch
# from .yolov4 import YoloBody
from .ssd import SSD
from .yolov3 import YoloV3


def get_model(opts):
    if opts.NAME == 'SSD':
        model = SSD(opts.NUM_CLASSES, opts.INPUT_SIZE, opts.BACKBONE, opts.BACKBONE_WEIGHTS)

    elif opts.NAME == 'YoloV3':
        model = YoloV3(opts.NUM_CLASSES, opts.INPUT_SIZE, opts.ANCHORS, opts.BACKBONE, opts.BACKBONE_WEIGHTS)
        
    elif opts.NAME == 'YoloV4':
        raise ValueError("model {} not supported yet".format(opts.NAME))
    elif opts.NAME == 'RetinaNet':
        raise ValueError("model {} not supported yet".format(opts.NAME))
    
    # load trained weights
    if opts.TRAINED_WEIGHTS is not None and os.path.exists(opts.TRAINED_WEIGHTS):
        state_dict = torch.load(opts.TRAINED_WEIGHTS, map_location='cpu')
        model.load_state_dict(state_dict)
    
    return model
